# RawServer.py
#
# Copyright (c) 2006 Michael Hobbs
#
# This file is part of iTorrent.
#
# iTorrent is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# iTorrent is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with iTorrent; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA


import socket
import struct

from BitTorrent.RawServer import RawServer, Handler


VERSION = 0x04
CONNECT = 0x01
NULL = 0x00
OK = 0x5a

# TODO: patch RawServer.start_connection
# TODO: patch RawServer_twisted
# TODO: patch Rerequester or zurlib

class patched_RawServer(RawServer):
	def async_start_connection(self, dns, handler=None, context=None, do_bind=True):
		proxy = self.config['socks_proxy']
		if proxy:
			handler = SocksHandler(dns, handler)
			host, port = proxy.split(':')
			dns = (host, int(port))
			context = None
		return super(patched_RawServer, self).async_start_connection(dns, handler, context, do_bind)

class SocksHandler(Handler):
	def __init__(self, dns, handler):
		super(SocksHandler, self).__init__()
		self.dns = dns
		self.handler = handler
		host, port = dns
		self.port = int(port)
		self.ip = socket.inet_aton(socket.gethostbyname(host))

	def connection_started(self, s):
		self.socket = s
		data = struct.pack('!BBH4sB', VERSION, CONNECT, self.port, self.ip, NULL)
		s.write(data)

	def connection_made(self, s):
		pass

	def connection_lost(self, s):
		s.handler = self.handler
		self.handler.connection_failed(self.dns, socket.error('Connection refused'))

	def connection_failed(self, addr, exception):
		self.handler.connection_failed(addr, exception)

	def connection_flushed(self, s):
		pass

	def data_came_in(self, unused_addr, datagram):
		null, status, port, ip = struct.unpack('!BBH4s', datagram)
		assert null == NULL
		assert port == self.port
		assert ip == self.ip
		self.socket.handler = self.handler
		if status == OK:
			self.handler.connection_started(self.socket)
		else:
			self.handler.connection_failed(self.dns, socket.error('Connection refused'))
		# end if
